from ruamel.yaml import YAML
from tensorflow.contrib.training import HParams
import os
import time


class YParams(HParams):

    def __init__(self, f_yaml, config_name):
        super().__init__()
        with open(f_yaml, 'r') as fp:
            for k, v in YAML().load(fp)[config_name].items():
                self.add_hparam(k, v)

    def save(self, file_name=None):
        if file_name is None:
            file_name = self.get('id_str') + '_hparams.yml'
        if not self.get('out_dir'):
            self.add_hparam('out_dir', '_'.join(
                (self.get('out_dir_prefix'), self.get('id_str'))))
        os.makedirs(self.get('out_dir'), exist_ok=True)
        file_path = os.path.join(self.get('out_dir'), file_name)
        with open(file_path, 'w') as outf:
            YAML().dump(self.values(), outf)


class HpsProcessor(object):
    def __init__(self, hparams, flags):
        self.hparams = hparams
        if flags is None:
            return
        for k, v in flags.items():
            if k in hparams:
                try:
                    hparams.set_hparam(k, v)
                except ValueError:
                    hparams.set_hparam(k, str(v))
            else:
                hparams.add_hparam(k, v)

    def _update_id_related(self):
        if not self.hparams.get('id_str'):
            id_str = time.strftime("%m%d%H%M", time.localtime())
            if self.hparams.get('id_prefix') and len(
                    self.hparams.get('id_prefix')) > 0:
                id_str = self.hparams.get('id_prefix') + '_' + id_str
            self.hparams.add_hparam('id_str', id_str)

    def _check_out_dir(self):
        if not self.hparams.get('out_dir'):
            if not self.get('out_dir'):
                self.hparams.add_hparam('out_dir', '_'.join(
                    (self.hparams.get('out_dir_prefix'),
                     self.hparams.get('id_str'))))
            os.makedirs(self.hparams.get('out_dir'), exist_ok=True)

    def _cuda_visiable_devices(self):
        if 'gpu' in self.hparams and self.hparams.gpu != '':
            if 'CUDA_VISIBLE_DEVICES' not in self.hparams:
                self.hparams.add_hparam('CUDA_VISIBLE_DEVICES',
                                        self.hparams.gpu)
            else:
                self.hparams.CUDA_VISIBLE_DEVICES = self.hparams.gpu
        if 'CUDA_VISIBLE_DEVICES' in self.hparams:
            os.environ[
                'CUDA_VISIBLE_DEVICES'] = self.hparams.CUDA_VISIBLE_DEVICES

    def process(self):
        self._update_id_related()
        self._check_out_dir()
        self._cuda_visiable_devices()
        return self.hparams
